"""
在gt_contour_point.py标注的基础上，再获取GT分割图像信息，用椭圆拟合边缘
"""

import numpy as np
import cv2
import os
import tqdm
import argparse
import yaml
from typing import Iterable, TextIO


def arg_parse():
    arg_parser = argparse.ArgumentParser()
    arg_parser.add_argument(
        "-c",
        "--collect_dir",
        type=str,
        help="The directory of the collected data",
    )
    arg_parser.add_argument(
        "--image", action="store_true", help="Whether to use image as background"
    )
    return arg_parser.parse_args()


class GTSegGen:
    def __init__(self, size, label_dir):
        with open(os.path.join(label_dir, "gt_labels.txt"), "r") as f:
            self.ind = map(lambda x: x.strip(), f.readlines())
        self.label_dir = label_dir
        # self.root_name = os.path.basename(label_dir)
        self.size = size
        self.output_dir = os.path.join(label_dir, "gt_segment")
        if not os.path.exists(self.output_dir):
            os.makedirs(self.output_dir)

    def __enter__(self):
        self.writer = open(os.path.join(self.label_dir, f"gt_segment.txt"), "w")
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.writer.close()
        # raise exc_type

    def __call__(self, use_bg: bool):
        for ind in tqdm.tqdm(self.ind, desc="Generating segment images"):
            # 对于每一张图像，读取标注点信息
            if use_bg:
                img = cv2.imread(os.path.join(self.label_dir, "images", f"{ind}.png"))
            else:
                img = np.zeros(self.size, dtype=np.uint8)
            with open(os.path.join(self.label_dir, "points_id", f"{ind}.txt")) as f:
                points_iter = self._read_label(f)
                img, th_num = self._fill_hole(points_iter, img)
            # 保存图像
            if th_num:
                print(ind, file=self.writer)
                cv2.imwrite(os.path.join(self.output_dir, f"{ind}.png"), img)

    def _read_label(self, ind: TextIO):
        """
        流式读取标注点的信息
        """
        _ = ind.readline()
        for i, line in enumerate(ind):
            x, y, idx = line.split(",")
            if i == 0:
                points = [(int(x), int(y))]
                idx_last = idx
            elif idx == idx_last:
                points.append((int(x), int(y)))
            else:
                yield points
                points = [(int(x), int(y))]
                idx_last = idx
        # 输出最后一个点列表
        yield points

    def _fill_hole(self, points_iter, mask):
        """
        用洪水淹没法填充孔洞，TODO 在边界上有部分孔洞无法填充
        """
        output = mask.copy()
        for points in map(np.array, points_iter):
            if points.shape[0] < 100:
                continue
            temp = np.zeros_like(mask, dtype=np.uint8)
            temp[points[:, 1], points[:, 0]] = 255
            temp = cv2.morphologyEx(
                temp,
                cv2.MORPH_CLOSE,
                cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5)),
            )
            cc_num, temp_mask = cv2.connectedComponents(temp)
            if cc_num > 2:
                continue
            cv2.floodFill(
                temp,
                np.pad(temp_mask.astype(np.uint8), 1, "constant", constant_values=1),
                (int(points[:, 0].mean()), int(points[:, 1].mean())),
                255,
            )
            output = cv2.bitwise_or(output, temp)
        if output.std() < 10 or (output != 0).sum() < 500:
            return output, False
        return output, True

    def _fit_ellipse(
        self, points_iter: Iterable, img: np.ndarray, tol=4, th_num=3
    ) -> np.ndarray:
        """
        拟合椭圆，TODO 但是在个别地方会出现误拟合的情况
        Arguments:
            points_iter (Iterable) : an iterator for points list with same id
            img (np.ndarray) : the image to fill ellipse
            tol (float) : the tolerance of the ratio of major axis and minor axis
            th_num (int) : the minimal number of points to fill ellipse
        """
        num_ellipse = 0
        for points in map(np.array, points_iter):
            if points.shape[0] > 10:  # 排除边缘点太少的陨石坑，避免误填充
                # 填充椭圆
                ellipse_box = cv2.fitEllipseAMS(points)
                # 长宽比例过大的椭圆不填充
                if (
                    ellipse_box[1][0] / ellipse_box[1][1] < 1 / tol
                    or ellipse_box[1][0] / ellipse_box[1][1] > tol
                ):
                    continue
                img = cv2.ellipse(img, ellipse_box, (255, 255, 255), -1)
                num_ellipse += 1

        return img, num_ellipse > th_num


if __name__ == "__main__":
    args = arg_parse()
    for collect_dir in args.collect_dir.split(","):
        with open(os.path.join(collect_dir, "config.yaml"), "r") as f:
            config = yaml.load(f.read(), Loader=yaml.FullLoader)
            size = (config["camera"]["height"], config["camera"]["width"])
        with GTSegGen(size, collect_dir) as gt_gen:
            gt_gen(args.image)
